import re
from langchain_core.output_parsers import BaseOutputParser


from langchain.prompts import (
    SystemMessagePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)

from llm.llm import llm


class jsonOutputParser(BaseOutputParser[str]):
    def parse(self, text: str) -> str:
        print("htmlOutputParser:", text)
        pattern = r"```json(.*?)```"
        matches = re.findall(pattern, text, re.DOTALL)
        if matches:
            return matches[0].strip()
        else:
            return "未找到HTML代码"

    @property
    def _type(self) -> str:
        return ""


def extract_data(desc, require=""):
    sys_tmp = """\
        角色设定：参数分析助手
        职责：首先分析输入数据，然后根据输出数据格式提取对应参数，最后根据输出数据格式返回结果。

        {{require}}

        输出数据：
        ```json
        [{{"label": "数据名称", "value": "数据数值"}}]
        ```
    """

    human_tmp = """\
        输入内容：
        {desc}
    """

    sys_msg = SystemMessagePromptTemplate.from_template(sys_tmp)
    human_msg = HumanMessagePromptTemplate.from_template(human_tmp)
    chat_msg = ChatPromptTemplate.from_messages([sys_msg, human_msg])

    data = ""
    try:
        chain = (chat_msg | llm | jsonOutputParser()).with_retry()
        data = chain.invoke(
            {"desc": desc, "require": require},
        )
        print("data", data)
    except Exception as e:
        print(e)
        return "数据分析失败，请重试"

    return data
